

<!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  
  <title>mindspore.ops.operations.custom_ops &mdash; MindSpore master documentation</title>
  

  
  <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />

  
  

  
  

  

  
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../../../" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/language_data.js"></script>
        <script async="async" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script>
    
    <script type="text/javascript" src="../../../../_static/js/theme.js"></script>

    
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../../../index.html" class="icon icon-home"> MindSpore
          

          
          </a>

          
            
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption"><span class="caption-text">MindSpore Python API</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.html">mindspore</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.common.initializer.html">mindspore.common.initializer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.communication.html">mindspore.communication</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.compression.html">mindspore.compression</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.context.html">mindspore.context</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.html">mindspore.dataset</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.audio.html">mindspore.dataset.audio</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.config.html">mindspore.dataset.config</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.text.html">mindspore.dataset.text</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.transforms.html">mindspore.dataset.transforms</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.vision.html">mindspore.dataset.vision</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.mindrecord.html">mindspore.mindrecord</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.html">mindspore.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.probability.html">mindspore.nn.probability</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.transformer.html">mindspore.nn.transformer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.numpy.html">mindspore.numpy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.ops.html">mindspore.ops</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.html">mindspore.parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.nn.html">mindspore.parallel.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.profiler.html">mindspore.profiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.scipy.html">mindspore.scipy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.train.html">mindspore.train</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.boost.html">mindspore.boost</a></li>
</ul>
<p class="caption"><span class="caption-text">MindSpore C++ API</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://www.mindspore.cn/lite/api/zh-CN/master/api_cpp/mindspore.html">MindSpore Lite↗</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">

      
      <nav class="wy-nav-top" aria-label="top navigation">
        
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">MindSpore</a>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        
          

















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
        
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
        
      <li>mindspore.ops.operations.custom_ops</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <h1>Source code for mindspore.ops.operations.custom_ops</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright 2021 Huawei Technologies Co., Ltd</span>
<span class="c1">#</span>
<span class="c1"># Licensed under the Apache License, Version 2.0 (the &quot;License&quot;);</span>
<span class="c1"># you may not use this file except in compliance with the License.</span>
<span class="c1"># You may obtain a copy of the License at</span>
<span class="c1">#</span>
<span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span>
<span class="c1">#</span>
<span class="c1"># Unless required by applicable law or agreed to in writing, software</span>
<span class="c1"># distributed under the License is distributed on an &quot;AS IS&quot; BASIS,</span>
<span class="c1"># WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.</span>
<span class="c1"># See the License for the specific language governing permissions and</span>
<span class="c1"># limitations under the License.</span>
<span class="c1"># ============================================================================</span>

<span class="sd">&quot;&quot;&quot;Custom operator&quot;&quot;&quot;</span>
<span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">inspect</span>
<span class="kn">import</span> <span class="nn">json</span>
<span class="kn">import</span> <span class="nn">hashlib</span>
<span class="kn">from</span> <span class="nn">mindspore</span> <span class="kn">import</span> <span class="n">ops</span>
<span class="kn">from</span> <span class="nn">mindspore</span> <span class="kn">import</span> <span class="n">log</span> <span class="k">as</span> <span class="n">logger</span>
<span class="kn">from</span> <span class="nn">mindspore.ops</span> <span class="kn">import</span> <span class="n">DataType</span>
<span class="kn">from</span> <span class="nn">mindspore.common</span> <span class="kn">import</span> <span class="n">dtype</span> <span class="k">as</span> <span class="n">mstype</span>
<span class="kn">from</span> <span class="nn">mindspore._c_expression</span> <span class="kn">import</span> <span class="n">Oplib</span>
<span class="kn">from</span> <span class="nn">._pyfunc_registry</span> <span class="kn">import</span> <span class="n">add_pyfunc</span>
<span class="kn">from</span> <span class="nn">._custom_grad</span> <span class="kn">import</span> <span class="n">autodiff_bprop</span>


<div class="viewcode-block" id="Custom"><a class="viewcode-back" href="../../../../api_python/ops/mindspore.ops.Custom.html#mindspore.ops.Custom">[docs]</a><span class="k">class</span> <span class="nc">Custom</span><span class="p">(</span><span class="n">ops</span><span class="o">.</span><span class="n">PrimitiveWithInfer</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    `Custom` primitive is used for user defined operators and is to enhance the expressive ability of built-in</span>
<span class="sd">    primitives. You can construct a `Custom` object with a predefined function, which describes the computation</span>
<span class="sd">    logic of a user defined operator. You can also construct another `Custom` object with another predefined</span>
<span class="sd">    function if needed. Then these `Custom` objects can be directly used in neural networks.</span>

<span class="sd">    .. warning::</span>
<span class="sd">        This is an experimental prototype that is subject to change.</span>

<span class="sd">    Args:</span>
<span class="sd">        func (Union[function, str]):</span>

<span class="sd">            - function: If func is of function type, then func should be a Python function which describes the</span>
<span class="sd">              computation logic of a user defined operator. The function can be one of the following:</span>

<span class="sd">              1. A AKG operator implementation function, which can use ir builder/tvm compute/hybrid grammar.</span>
<span class="sd">              2. A TBE operator implementation function.</span>
<span class="sd">              3. A pure python function</span>

<span class="sd">            - str: If func is of str type, then str should be a path of binary file along with a function name.</span>
<span class="sd">              This could only be used when func_type is &quot;aot&quot;. Currently &quot;aot&quot; supports GPU/CPU(linux only) platform.</span>
<span class="sd">              &quot;aot&quot; means ahead of time, in which case Custom directly launches user defined &quot;xxx.so&quot; file as an</span>
<span class="sd">              operator. Users need to compile a handwriting &quot;xxx.cu&quot;/&quot;xxx.cc&quot; file into &quot;xxx.so&quot; ahead of time,</span>
<span class="sd">              and offer the path of the file along with a function name.</span>

<span class="sd">              - &quot;xxx.so&quot; file generation:</span>

<span class="sd">                1) GPU Platform: Given user defined &quot;xxx.cu&quot; file (ex. &quot;{path}/add.cu&quot;), use nvcc command to compile</span>
<span class="sd">                it.(ex. &quot;nvcc --shared -Xcompiler -fPIC -o add.so add.cu&quot;)</span>

<span class="sd">                2) CPU Platform: Given user defined &quot;xxx.cc&quot; file (ex. &quot;{path}/add.cc&quot;), use g++/gcc command to compile</span>
<span class="sd">                it.(ex. &quot;g++ --shared -fPIC  -o add.so add.cc&quot;)</span>

<span class="sd">              - Define a &quot;xxx.cc&quot;/&quot;xxx.cu&quot; file:</span>

<span class="sd">                &quot;aot&quot; is a cross-platform identity. The functions defined in &quot;xxx.cc&quot; or &quot;xxx.cu&quot; share the same args.</span>
<span class="sd">                Typically, the function should be as:</span>

<span class="sd">                .. code-block::</span>

<span class="sd">                    int func(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes,</span>
<span class="sd">                             void *stream, void *extra)</span>

<span class="sd">                Parameters:</span>

<span class="sd">                - nparam(int): total number of inputs plus outputs; suppose the operator has 2 inputs and 3 outputs,</span>
<span class="sd">                  then nparam=5</span>
<span class="sd">                - params(void \*\*): a pointer to the array of inputs and outputs&#39; pointer; the pointer type of inputs</span>
<span class="sd">                  and outputs is void \* ; suppose the operator has 2 inputs and 3 outputs, then the first input&#39;s</span>
<span class="sd">                  pointer is params[0] and the second output&#39;s pointer is params[3]</span>
<span class="sd">                - ndims(int \*): a pointer to the array of inputs and outputs&#39; dimension num; suppose params[i] is a</span>
<span class="sd">                  1024x1024 tensor and params[j] is a 77x83x4 tensor, then ndims[i]=2, ndims[j]=3.</span>
<span class="sd">                - shapes(int64_t \*\*): a pointer to the array of inputs and outputs&#39; shapes(int64_t \*); the ith</span>
<span class="sd">                  input&#39;s jth dimension&#39;s size is shapes[i][j](0&lt;=j&lt;ndims[i]); suppose params[i] is a 2x3 tensor and</span>
<span class="sd">                  params[j] is a 3x3x4 tensor, then shapes[i][0]=2, shapes[j][2]=4.</span>
<span class="sd">                - dtypes(const char \*\*): a pointer to the array of inputs and outputs&#39; types(const char \*);</span>
<span class="sd">                  (ex. &quot;float32&quot;, &quot;float16&quot;, &quot;float&quot;, &quot;float64&quot;, &quot;int&quot;, &quot;int8&quot;, &quot;int16&quot;, &quot;int32&quot;, &quot;int64&quot;, &quot;uint&quot;,</span>
<span class="sd">                  &quot;uint8&quot;, &quot;uint16&quot;, &quot;uint32&quot;, &quot;uint64&quot;, &quot;bool&quot;)</span>
<span class="sd">                - stream(void \*): stream pointer, only used in cuda file</span>
<span class="sd">                - extra(void \*): used for further extension</span>

<span class="sd">                Return Value(int):</span>

<span class="sd">                - 0: raise no Exception</span>
<span class="sd">                - larger than 0: will raise Exception</span>

<span class="sd">                Examples: see details tests/st/ops/graph_kernel/custom/aot_test_files/</span>

<span class="sd">              - Use it in Custom:</span>

<span class="sd">                .. code-block::</span>

<span class="sd">                    Custom(func=&quot;{path}/{file_name}:{func_name}&quot;,...)</span>
<span class="sd">                    (ex. Custom(func=&quot;./reorganize.so:CustomReorganize&quot;, out_shape=[1], out_dtype=mstype.float32))</span>

<span class="sd">        out_shape (Union[function, list, tuple]): The output shape infer function or the value of output shape of</span>
<span class="sd">            `func`.</span>

<span class="sd">            If func has single output, then the value of output shape is a list or tuple of int.</span>

<span class="sd">            If func has multiple outputs, then the value of output shape is a tuple, each item represents the shape</span>
<span class="sd">            of each output.</span>

<span class="sd">        out_dtype (Union[function, :class:`mindspore.dtype`, tuple[:class:`mindspore.dtype`]]): The output data type</span>
<span class="sd">            infer function or the value of output data type of `func`.</span>

<span class="sd">            If func has single output, then the value of output shape is a `mindspore.dtype`.</span>

<span class="sd">            If func has multiple outputs, then the value of output shape is a tuple of `mindspore.dtype`, each item</span>
<span class="sd">            represents the data type of each output.</span>

<span class="sd">        func_type (str): The implementation type of `func`, should be one of [&quot;akg&quot;, &quot;tbe&quot;, &quot;aot&quot;, &quot;pyfunc&quot;]. Each</span>
<span class="sd">            `func_type` only supports specific platforms(targets). The supported platforms of `func_type`:</span>

<span class="sd">            - &quot;akg&quot;: supports [&quot;Ascend&quot;, &quot;GPU&quot;].</span>
<span class="sd">            - &quot;tbe&quot;: supports [&quot;Ascend&quot;].</span>
<span class="sd">            - &quot;aot&quot;: supports [&quot;GPU&quot;, &quot;CPU&quot;].</span>
<span class="sd">            - &quot;pyfunc&quot;: supports [&quot;CPU&quot;].</span>

<span class="sd">        bprop (function): The back propagation function of `func`. Default: None.</span>
<span class="sd">        reg_info (Union[str, dict, list, tuple]): Represents the registration information(reg info) of `func` with</span>
<span class="sd">            json format of type str or dict. The reg info specifies supported data types and formats of inputs and</span>
<span class="sd">            outputs, attributes and target of `func`. Default: None.</span>

<span class="sd">            If reg info is a list or tuple, then each item should be with json format of type str or dict, which</span>
<span class="sd">            represents the registration information of `func` in a specific target. You need to invoke `CustomRegOp`</span>
<span class="sd">            or the subclass of `RegOp` to generate the reg info for `func`. Then you can invoke</span>
<span class="sd">            `custom_info_register` to bind the reg info to `func` or just pass the reg info to `reg_info` parameter.</span>
<span class="sd">            The `reg_info` parameter takes higher priority than `custom_info_register` and the reg info in a</span>
<span class="sd">            specific target will be registered only once.</span>

<span class="sd">            If reg info is not set, then we will infer the data types and formats from the inputs of `Custom` operator.</span>

<span class="sd">            Please note that, if `func_type` is &quot;tbe&quot; or the `func` only supports some specified data types and formats,</span>
<span class="sd">            or it has attribute inputs, then you should set the reg info for `func`.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **input** (Union(tuple, list)) - The input tuple or list is made up of multiple tensors, and attributes</span>
<span class="sd">          value(optional).</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor or tuple[Tensor], execution results.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If the type of `func` is invalid or the type of register information for `func` is invalid.</span>
<span class="sd">        ValueError: If `func_type` is invalid.</span>
<span class="sd">        ValueError: If the register information is invalid, including the target is not supported, the input numbers</span>
<span class="sd">            or the attributes of `func` differs in different targets.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; import mindspore.ops as ops</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.ops import CustomRegOp, custom_info_register, DataType</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.common import dtype as mstype</span>
<span class="sd">        &gt;&gt;&gt; from mindspore.nn import Cell</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; # Example, func_type = &quot;akg&quot;</span>
<span class="sd">        &gt;&gt;&gt; def outer_product(a, b):</span>
<span class="sd">        ...     c = output_tensor(a.shape, a.dtype)</span>
<span class="sd">        ...     for i0 in range(a.shape[0]):</span>
<span class="sd">        ...         for i1 in range(b.shape[1]):</span>
<span class="sd">        ...             c[i0, i1] = 0.0</span>
<span class="sd">        ...             for i2 in range(a.shape[1]):</span>
<span class="sd">        ...                 c[i0, i1] = c[i0, i1] + (a[i0, i2] * b[i2, i1])</span>
<span class="sd">        ...     return c</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; class AkgNet(Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(AkgNet, self).__init__()</span>
<span class="sd">        ...         def infer_func(x, y):</span>
<span class="sd">        ...             return x</span>
<span class="sd">        ...         self.program = ops.Custom(outer_product, out_shape=infer_func, out_dtype=infer_func, \</span>
<span class="sd">        ...                                   func_type=&quot;akg&quot;)</span>
<span class="sd">        ...     def construct(self, x, y):</span>
<span class="sd">        ...         return self.program(x, y)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; # Example, func_type = &quot;tbe&quot;</span>
<span class="sd">        &gt;&gt;&gt; square_with_bias_op_info = CustomRegOp() \</span>
<span class="sd">        ...     .fusion_type(&quot;OPAQUE&quot;) \</span>
<span class="sd">        ...     .attr(&quot;bias&quot;, &quot;required&quot;, &quot;float&quot;) \</span>
<span class="sd">        ...     .input(0, &quot;x&quot;) \</span>
<span class="sd">        ...     .output(0, &quot;y&quot;) \</span>
<span class="sd">        ...     .dtype_format(DataType.F32_Default, DataType.F32_Default) \</span>
<span class="sd">        ...     .dtype_format(DataType.F16_Default, DataType.F16_Default) \</span>
<span class="sd">        ...     .target(&quot;Ascend&quot;) \</span>
<span class="sd">        ...     .get_op_info()</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; @custom_info_register(square_with_bias_op_info)</span>
<span class="sd">        ... def square_with_bias(input_x, output_y, bias=0.0, kernel_name=&quot;square_with_bias&quot;):</span>
<span class="sd">        ...     import te.lang.cce</span>
<span class="sd">        ...     from te import tvm</span>
<span class="sd">        ...     from topi.cce import util</span>
<span class="sd">        ...</span>
<span class="sd">        ...     shape = input_x.get(&quot;shape&quot;)</span>
<span class="sd">        ...     dtype = input_x.get(&quot;dtype&quot;).lower()</span>
<span class="sd">        ...</span>
<span class="sd">        ...     shape = util.shape_refine(shape)</span>
<span class="sd">        ...     data = tvm.placeholder(shape, name=&quot;data&quot;, dtype=dtype)</span>
<span class="sd">        ...</span>
<span class="sd">        ...     with tvm.target.cce():</span>
<span class="sd">        ...         res0 = te.lang.cce.vmul(data, data)</span>
<span class="sd">        ...         res = te.lang.cce.vadds(res0, bias)</span>
<span class="sd">        ...         sch = te.lang.cce.auto_schedule(res)</span>
<span class="sd">        ...</span>
<span class="sd">        ...     config = {&quot;print_ir&quot;: False,</span>
<span class="sd">        ...               &quot;name&quot;: kernel_name,</span>
<span class="sd">        ...               &quot;tensor_list&quot;: [data, res]}</span>
<span class="sd">        ...</span>
<span class="sd">        ...     te.lang.cce.cce_build_code(sch, config)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; class TbeNet(Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(TbeNet, self).__init__()</span>
<span class="sd">        ...         self.square_with_bias = ops.Custom(square_with_bias, out_shape=lambda x, _: x, \</span>
<span class="sd">        ...                                            out_dtype=lambda x, _: x, func_type=&quot;tbe&quot;)</span>
<span class="sd">        ...     def construct(self, x):</span>
<span class="sd">        ...         res = self.square_with_bias(x, 1.0)</span>
<span class="sd">        ...         return res</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; # Example, func_type = &quot;aicpu&quot;</span>
<span class="sd">        &gt;&gt;&gt; resize_bilinear_op_info = CustomRegOp(&quot;ResizeBilinear&quot;) \</span>
<span class="sd">        ...     .fusion_type(&quot;OPAQUE&quot;) \</span>
<span class="sd">        ...     .input(0, &quot;input&quot;, &quot;required&quot;) \</span>
<span class="sd">        ...     .output(1, &quot;output&quot;, &quot;required&quot;) \</span>
<span class="sd">        ...     .attr(&quot;align_corners&quot;, &quot;required&quot;, &quot;bool&quot;) \</span>
<span class="sd">        ...     .attr(&quot;cust_aicpu&quot;, &quot;optional&quot;, &quot;str&quot;, &quot;aicpu_kernels&quot;) \</span>
<span class="sd">        ...     .dtype_format(DataType.F32_Default, DataType.F32_Default) \</span>
<span class="sd">        ...     .dtype_format(DataType.F16_Default, DataType.F32_Default) \</span>
<span class="sd">        ...     .target(&quot;Ascend&quot;) \</span>
<span class="sd">        ...     .get_op_info()</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; @custom_info_register(resize_bilinear_op_info)</span>
<span class="sd">        ... def resize_bilinear_aicpu():</span>
<span class="sd">        ...     return</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; class AicpuNet(Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(AicpuNet, self).__init__()</span>
<span class="sd">        ...         self.resize_bilinear_op = ops.Custom(resize_bilinear_aicpu, out_shape=[1, 1, 9, 9], \</span>
<span class="sd">        ...                                              out_dtype=mstype.float32, func_type=&quot;aicpu&quot;)</span>
<span class="sd">        ...     def construct(self, x):</span>
<span class="sd">        ...         res = self.resize_bilinear_op(x, True, &quot;aicpu_kernels&quot;)</span>
<span class="sd">        ...         return res</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; # Example, func_type = &quot;aot&quot;</span>
<span class="sd">        &gt;&gt;&gt; class AOTSingleOutputNet(Cell):</span>
<span class="sd">        ...     def __init__(self, out_shapes, out_types):</span>
<span class="sd">        ...         super(AOTSingleOutputNet, self).__init__()</span>
<span class="sd">        ...         self.program = ops.Custom(&quot;./reorganize.so:CustomReorganize&quot;, out_shapes, out_types, &quot;aot&quot;)</span>
<span class="sd">        ...     def construct(self, x, y):</span>
<span class="sd">        ...         return self.program(x, y)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; # Example, func_type = &quot;pyfunc&quot;</span>
<span class="sd">        &gt;&gt;&gt; def func_multi_output(x1, x2):</span>
<span class="sd">        ...     return (x1 + x2), (x1 - x2)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; class PyFuncNet(Cell):</span>
<span class="sd">        ...     def __init__(self):</span>
<span class="sd">        ...         super(PyFuncNet, self).__init__()</span>
<span class="sd">        ...         self.func = ops.Custom(func_multi_output, lambda x, _: (x, x), lambda x, _: (x, x), &quot;pyfunc&quot;)</span>
<span class="sd">        ...     def construct(self, x1, x2):</span>
<span class="sd">        ...         return self.func(x1, x2)</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">registered_func</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="n">attr_dict</span> <span class="o">=</span> <span class="p">{}</span>  <span class="c1"># Save input_names and attr_names for func.</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">func</span><span class="p">,</span> <span class="n">out_shape</span><span class="p">,</span> <span class="n">out_dtype</span><span class="p">,</span> <span class="n">func_type</span><span class="p">,</span> <span class="n">bprop</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">reg_info</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="n">ops</span><span class="o">.</span><span class="n">PrimitiveWithInfer</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s2">&quot;Custom&quot;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">supported_targets</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;Ascend&quot;</span><span class="p">,</span> <span class="s2">&quot;GPU&quot;</span><span class="p">,</span> <span class="s2">&quot;CPU&quot;</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">supported_func_type</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;akg&quot;</span><span class="p">,</span> <span class="s2">&quot;tbe&quot;</span><span class="p">,</span> <span class="s2">&quot;aicpu&quot;</span><span class="p">,</span> <span class="s2">&quot;aot&quot;</span><span class="p">,</span> <span class="s2">&quot;pyfunc&quot;</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">func</span> <span class="o">=</span> <span class="n">func</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span> <span class="o">=</span> <span class="n">func_type</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">func_name</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">uniq_name</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">imply_path</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">func_source_str</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_check_func</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_update_func_info</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;func_name&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;uniq_name&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">uniq_name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;imply_path&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">imply_path</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span> <span class="o">==</span> <span class="s2">&quot;pyfunc&quot;</span><span class="p">:</span>
            <span class="n">func_id</span> <span class="o">=</span> <span class="nb">id</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">)</span>
            <span class="n">add_pyfunc</span><span class="p">(</span><span class="n">func_id</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;fn_id&quot;</span><span class="p">,</span> <span class="n">func_id</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">out_shape</span> <span class="o">=</span> <span class="n">out_shape</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">out_dtype</span> <span class="o">=</span> <span class="n">out_dtype</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">bprop</span> <span class="o">=</span> <span class="n">bprop</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fake_output</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">single_scalar_output</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">out_dtype</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">fake_output</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="k">elif</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">out_shape</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">single_scalar_output</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;fake_output&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">fake_output</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;single_scalar_output&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">single_scalar_output</span><span class="p">)</span>

        <span class="c1"># Register info</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_register_info</span><span class="p">(</span><span class="n">reg_info</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">func_type</span> <span class="o">==</span> <span class="s2">&quot;akg&quot;</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s1">&#39;func_source_str&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_source_str</span><span class="p">)</span>
            <span class="k">if</span> <span class="s2">&quot;ir_builder&quot;</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_source_str</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span> <span class="o">=</span> <span class="s2">&quot;ir_builder&quot;</span>
            <span class="k">elif</span> <span class="s2">&quot;compute&quot;</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_source_str</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span> <span class="o">=</span> <span class="s2">&quot;tvm_compute&quot;</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span> <span class="o">=</span> <span class="s2">&quot;hybrid&quot;</span>
                <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">bprop</span><span class="p">:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">_hybrid_autodiff</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;func_type&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_update_attr</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">infer_shape</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">callable</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">out_shape</span><span class="p">):</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">out_shape</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">out_shape</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">out_shape</span>
        <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;The function output are empty tuple. Add a placeholder instead. &quot;</span>
                       <span class="s2">&quot;Do not use it as it could be any uninitialized data.&quot;</span><span class="p">)</span>
        <span class="k">return</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)</span>

    <span class="k">def</span> <span class="nf">infer_dtype</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">callable</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">out_dtype</span><span class="p">):</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">out_dtype</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">out_dtype</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">out_dtype</span>
        <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;The function output are empty tuple. Add a placeholder instead. &quot;</span>
                       <span class="s2">&quot;Do not use it as it could be any uninitialized data.&quot;</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">mstype</span><span class="o">.</span><span class="n">int32</span>

    <span class="k">def</span> <span class="nf">get_bprop</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">bprop</span>

    <span class="k">def</span> <span class="nf">_check_func</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Check the validity of func_type and type of func&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span> <span class="ow">not</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">supported_func_type</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;func_type should be one of </span><span class="si">{}</span><span class="s2">, but got </span><span class="si">{}</span><span class="s2">&quot;</span>
                             <span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">supported_func_type</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span><span class="p">))</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span> <span class="o">==</span> <span class="s2">&quot;aot&quot;</span><span class="p">:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
                <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;</span><span class="si">{}</span><span class="s2"> func should be of type str, but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func_type</span><span class="p">,</span> <span class="nb">type</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">)))</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">callable</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">):</span>
                <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;</span><span class="si">{}</span><span class="s2"> func should be of type function, but got </span><span class="si">{}</span><span class="s2">&quot;</span>
                                <span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func_type</span><span class="p">,</span> <span class="nb">type</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">)))</span>

    <span class="k">def</span> <span class="nf">_update_func_info</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Update information of func&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">callable</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">):</span>
            <span class="c1"># Get the original function if func is decorated</span>
            <span class="k">if</span> <span class="s2">&quot;__wrapped__&quot;</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">func</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">[</span><span class="s2">&quot;__wrapped__&quot;</span><span class="p">]</span>
            <span class="c1"># func name</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">func_name</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="o">.</span><span class="vm">__name__</span>
            <span class="c1"># path of func</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">imply_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">realpath</span><span class="p">(</span><span class="n">inspect</span><span class="o">.</span><span class="n">getfile</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">))</span>
            <span class="c1"># source code of func, not include the decorator before def</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">func_source_str</span> <span class="o">=</span> <span class="n">inspect</span><span class="o">.</span><span class="n">getsource</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">)</span>
            <span class="n">index</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_source_str</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="s2">&quot;def &quot;</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">index</span> <span class="o">!=</span> <span class="o">-</span><span class="mi">1</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">func_source_str</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_source_str</span><span class="p">[</span><span class="n">index</span><span class="p">:]</span>
            <span class="c1"># unique func name</span>
            <span class="n">sha256</span> <span class="o">=</span> <span class="n">hashlib</span><span class="o">.</span><span class="n">sha256</span><span class="p">()</span>
            <span class="n">sha256</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">imply_path</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s2">&quot;utf-8&quot;</span><span class="p">))</span>
            <span class="n">sha256</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func_source_str</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s2">&quot;utf-8&quot;</span><span class="p">))</span>
            <span class="n">hash_str</span> <span class="o">=</span> <span class="n">sha256</span><span class="o">.</span><span class="n">hexdigest</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">uniq_name</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span> <span class="o">+</span> <span class="s2">&quot;_&quot;</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_name</span> <span class="o">+</span> <span class="s2">&quot;_&quot;</span> <span class="o">+</span> <span class="n">hash_str</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="c1"># func name</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">func_name</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">func</span>
            <span class="c1"># uniq func name</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">uniq_name</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span> <span class="o">+</span> <span class="s2">&quot;_&quot;</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_name</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;func should be of type function or str, but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">)))</span>

    <span class="k">def</span> <span class="nf">_register_info</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">info</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Register reg_info.&quot;&quot;&quot;</span>
        <span class="n">reg_info</span> <span class="o">=</span> <span class="n">info</span>
        <span class="k">if</span> <span class="n">reg_info</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="s2">&quot;reg_info&quot;</span><span class="p">):</span>
            <span class="n">reg_info</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="s2">&quot;reg_info&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span> <span class="o">==</span> <span class="s2">&quot;aicpu&quot;</span> <span class="ow">and</span> <span class="n">reg_info</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;custom aicpu ops must set reg_info, but current reg_info is None.&quot;</span><span class="p">)</span>
        <span class="n">reg_info_list</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_expanded_list</span><span class="p">(</span><span class="n">reg_info</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">reg_info</span> <span class="ow">in</span> <span class="n">reg_info_list</span><span class="p">:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">reg_info</span><span class="p">,</span> <span class="p">(</span><span class="nb">str</span><span class="p">,</span> <span class="nb">dict</span><span class="p">)):</span>
                <span class="k">continue</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">reg_info</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
                <span class="n">reg_info</span> <span class="o">=</span> <span class="n">json</span><span class="o">.</span><span class="n">loads</span><span class="p">(</span><span class="n">reg_info</span><span class="p">)</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">fake_output</span><span class="p">:</span>
                <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;outputs&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">dict</span><span class="p">({</span><span class="s2">&quot;index&quot;</span><span class="p">:</span> <span class="mi">0</span><span class="p">,</span> <span class="s2">&quot;name&quot;</span><span class="p">:</span> <span class="s2">&quot;y&quot;</span><span class="p">,</span> <span class="s2">&quot;param_type&quot;</span><span class="p">:</span> <span class="s2">&quot;required&quot;</span><span class="p">}))</span>
                <span class="n">new_dtype_format</span> <span class="o">=</span> <span class="p">[]</span>
                <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;dtype_format&quot;</span><span class="p">]:</span>
                    <span class="n">new_dtype_format</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">i</span> <span class="o">+</span> <span class="p">(</span><span class="n">DataType</span><span class="o">.</span><span class="n">I32_Default</span><span class="p">,))</span>
                <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;dtype_format&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">new_dtype_format</span>
            <span class="n">target</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_target</span><span class="p">(</span><span class="n">reg_info</span><span class="p">)</span>
            <span class="c1"># Reg info for func is only registered once for a certain target</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_has_registered</span><span class="p">(</span><span class="n">target</span><span class="p">):</span>
                <span class="k">continue</span>
            <span class="c1"># Register</span>
            <span class="n">reg_info</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_reformat_reg_info</span><span class="p">(</span><span class="n">reg_info</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
            <span class="n">reg_info_str</span> <span class="o">=</span> <span class="n">json</span><span class="o">.</span><span class="n">dumps</span><span class="p">(</span><span class="n">reg_info</span><span class="p">)</span>
            <span class="n">op_lib</span> <span class="o">=</span> <span class="n">Oplib</span><span class="p">()</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">op_lib</span><span class="o">.</span><span class="n">reg_op</span><span class="p">(</span><span class="n">reg_info_str</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">imply_path</span><span class="p">):</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;Invalid reg info </span><span class="si">{}</span><span class="s1">: </span><span class="si">{}</span><span class="se">\n</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">imply_path</span><span class="p">,</span> <span class="n">reg_info_str</span><span class="p">))</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_save_attr</span><span class="p">(</span><span class="n">reg_info</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_save_register_status</span><span class="p">(</span><span class="n">target</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_get_expanded_list</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">data</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Recursive function to parse elements in list or tuple.&quot;&quot;&quot;</span>
        <span class="n">data_list</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">)):</span>
            <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">data</span><span class="p">:</span>
                <span class="n">tmp_list</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_expanded_list</span><span class="p">(</span><span class="n">i</span><span class="p">)</span>
                <span class="k">for</span> <span class="n">ii</span> <span class="ow">in</span> <span class="n">tmp_list</span><span class="p">:</span>
                    <span class="n">data_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">ii</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">data_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">data_list</span>

    <span class="k">def</span> <span class="nf">_get_registered_targets</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Get the registered targets of func.&quot;&quot;&quot;</span>
        <span class="n">targets</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">if</span> <span class="n">callable</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">):</span>
            <span class="n">targets</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="s2">&quot;registered_targets&quot;</span><span class="p">,</span> <span class="p">[])</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="n">targets</span> <span class="o">=</span> <span class="n">Custom</span><span class="o">.</span><span class="n">registered_func</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="p">[])</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">targets</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="n">targets</span> <span class="o">=</span> <span class="p">[</span><span class="n">targets</span><span class="p">]</span>
        <span class="k">return</span> <span class="n">targets</span>

    <span class="k">def</span> <span class="nf">_has_registered</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Check if registration information is registered in target.&quot;&quot;&quot;</span>
        <span class="n">registered_targets</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_registered_targets</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">target</span> <span class="ow">in</span> <span class="n">registered_targets</span>

    <span class="k">def</span> <span class="nf">_save_register_status</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Save registration status for target.&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">callable</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">):</span>
            <span class="n">registered_targets</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="s2">&quot;registered_targets&quot;</span><span class="p">,</span> <span class="p">[])</span>
            <span class="n">registered_targets</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">target</span><span class="p">)</span>
            <span class="nb">setattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="s2">&quot;registered_targets&quot;</span><span class="p">,</span> <span class="n">registered_targets</span><span class="p">)</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">Custom</span><span class="o">.</span><span class="n">registered_func</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">),</span> <span class="nb">list</span><span class="p">):</span>
                <span class="n">Custom</span><span class="o">.</span><span class="n">registered_func</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">target</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">Custom</span><span class="o">.</span><span class="n">registered_func</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span><span class="n">target</span><span class="p">]</span>

    <span class="k">def</span> <span class="nf">_get_op_name</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">reg_info</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span> <span class="o">==</span> <span class="s2">&quot;aicpu&quot;</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">uniq_name</span> <span class="o">=</span> <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;op_name&quot;</span><span class="p">]</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;uniq_name&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">uniq_name</span><span class="p">)</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">uniq_name</span>

    <span class="k">def</span> <span class="nf">_reformat_reg_info</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">reg_info</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Reformat registration information.&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">reg_info</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;reg_info should be of type dict, but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">reg_info</span><span class="p">)))</span>
        <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;op_name&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_op_name</span><span class="p">(</span><span class="n">reg_info</span><span class="p">)</span>
        <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;imply_type&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_imply_type</span><span class="p">(</span><span class="n">reg_info</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;fusion_type&quot;</span><span class="p">),</span> <span class="nb">str</span><span class="p">)</span> <span class="ow">or</span> <span class="ow">not</span> <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;fusion_type&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">strip</span><span class="p">():</span>
            <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;fusion_type&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;OPAQUE&quot;</span>
        <span class="c1"># Supplement necessary info for TBE if these information is missing in reg_info</span>
        <span class="k">if</span> <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;imply_type&quot;</span><span class="p">]</span> <span class="o">==</span> <span class="s2">&quot;TBE&quot;</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;attr&quot;</span><span class="p">)</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;attr&quot;</span><span class="p">],</span> <span class="nb">list</span><span class="p">):</span>
                <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">item</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;attr&quot;</span><span class="p">]):</span>
                    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">item</span><span class="p">,</span> <span class="nb">dict</span><span class="p">)</span> <span class="ow">and</span> <span class="n">item</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;value&quot;</span><span class="p">)</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                        <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;attr&quot;</span><span class="p">][</span><span class="n">i</span><span class="p">][</span><span class="s2">&quot;value&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;all&quot;</span>
            <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;async_flag&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;async_flag&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
            <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;binfile_name&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_name</span> <span class="o">+</span> <span class="s2">&quot;.so&quot;</span>
            <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;compute_cost&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;compute_cost&quot;</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>
            <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;kernel_name&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_name</span>
            <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;partial_flag&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;partial_flag&quot;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
            <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;need_check_supported&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;need_check_supported&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
        <span class="c1"># Supplement necessary info for AKG if these information is missing in reg_info</span>
        <span class="k">if</span> <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;imply_type&quot;</span><span class="p">]</span> <span class="o">==</span> <span class="s2">&quot;AKG&quot;</span><span class="p">:</span>
            <span class="n">target_to_processor</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;Ascend&quot;</span><span class="p">:</span> <span class="s2">&quot;AiCore&quot;</span><span class="p">,</span> <span class="s2">&quot;GPU&quot;</span><span class="p">:</span> <span class="s2">&quot;CUDA&quot;</span><span class="p">,</span> <span class="s2">&quot;CPU&quot;</span><span class="p">:</span> <span class="s2">&quot;CPU&quot;</span><span class="p">}</span>
            <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;processor&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;processor&quot;</span><span class="p">,</span> <span class="n">target_to_processor</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">target</span><span class="p">))</span>
        <span class="k">return</span> <span class="n">reg_info</span>

    <span class="k">def</span> <span class="nf">_get_target</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">reg_info</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Get target information.&quot;&quot;&quot;</span>
        <span class="n">target</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">reg_info</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
            <span class="c1"># Get target from reg_info[&quot;target&quot;]</span>
            <span class="n">target</span> <span class="o">=</span> <span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;target&quot;</span><span class="p">)</span>
            <span class="c1"># Infer target from reg_info[&quot;processor&quot;], reg_info generated from AkgGpuRegOp or AkgAscendRegOp</span>
            <span class="c1">#   will have the processor information.</span>
            <span class="k">if</span> <span class="n">target</span> <span class="ow">not</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">supported_targets</span><span class="p">:</span>
                <span class="n">processor_to_target</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;AiCore&quot;</span><span class="p">:</span> <span class="s2">&quot;Ascend&quot;</span><span class="p">,</span> <span class="s2">&quot;CUDA&quot;</span><span class="p">:</span> <span class="s2">&quot;GPU&quot;</span><span class="p">,</span> <span class="s2">&quot;CPU&quot;</span><span class="p">:</span> <span class="s2">&quot;CPU&quot;</span><span class="p">}</span>
                <span class="n">target</span> <span class="o">=</span> <span class="n">processor_to_target</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;processor&quot;</span><span class="p">))</span>
            <span class="c1"># Infer target from reg_info[&quot;imply_type&quot;]</span>
            <span class="k">if</span> <span class="n">target</span> <span class="ow">not</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">supported_targets</span><span class="p">:</span>
                <span class="n">imply_type_to_target</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;TBE&quot;</span><span class="p">:</span> <span class="s2">&quot;Ascend&quot;</span><span class="p">,</span> <span class="s2">&quot;GPU&quot;</span><span class="p">:</span> <span class="s2">&quot;GPU&quot;</span><span class="p">,</span> <span class="s2">&quot;CPU&quot;</span><span class="p">:</span> <span class="s2">&quot;CPU&quot;</span><span class="p">}</span>
                <span class="n">target</span> <span class="o">=</span> <span class="n">imply_type_to_target</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;imply_type&quot;</span><span class="p">))</span>
        <span class="c1"># Infer target from func_type</span>
        <span class="k">if</span> <span class="n">target</span> <span class="ow">not</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">supported_targets</span><span class="p">:</span>
            <span class="n">func_type_to_target</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;tbe&quot;</span><span class="p">:</span> <span class="s2">&quot;Ascend&quot;</span><span class="p">,</span> <span class="s2">&quot;pyfunc&quot;</span><span class="p">:</span> <span class="s2">&quot;CPU&quot;</span><span class="p">}</span>
            <span class="n">target</span> <span class="o">=</span> <span class="n">func_type_to_target</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func_type</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">target</span> <span class="ow">not</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">supported_targets</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;target should be one of </span><span class="si">{}</span><span class="s2">, but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">supported_targets</span><span class="p">,</span> <span class="n">target</span><span class="p">))</span>
        <span class="k">return</span> <span class="n">target</span>

    <span class="k">def</span> <span class="nf">_get_imply_type</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">reg_info</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Get imply_typ information.&quot;&quot;&quot;</span>
        <span class="c1"># Get imply_type from reg_info[&quot;imply_type&quot;]</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">reg_info</span><span class="p">,</span> <span class="nb">dict</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;imply_type&quot;</span><span class="p">),</span> <span class="nb">str</span><span class="p">)</span> <span class="ow">and</span> \
                <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;imply_type&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">strip</span><span class="p">():</span>
            <span class="k">return</span> <span class="n">reg_info</span><span class="p">[</span><span class="s2">&quot;imply_type&quot;</span><span class="p">]</span>
        <span class="c1"># Infer imply_type from func_type</span>
        <span class="n">func_type_to_imply_type</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;akg&quot;</span><span class="p">:</span> <span class="s2">&quot;AKG&quot;</span><span class="p">,</span> <span class="s2">&quot;tbe&quot;</span><span class="p">:</span> <span class="s2">&quot;TBE&quot;</span><span class="p">,</span> <span class="s2">&quot;aicpu&quot;</span><span class="p">:</span> <span class="s2">&quot;AiCPU&quot;</span><span class="p">,</span> <span class="s2">&quot;aot&quot;</span><span class="p">:</span> <span class="n">target</span><span class="p">,</span> <span class="s2">&quot;pyfunc&quot;</span><span class="p">:</span> <span class="n">target</span><span class="p">}</span>
        <span class="k">return</span> <span class="n">func_type_to_imply_type</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func_type</span><span class="p">,</span> <span class="s2">&quot;AKG&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_save_attr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">reg_info</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Save input_names and attr_names of current func.&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">reg_info</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
            <span class="k">return</span>
        <span class="n">tensor_inputs</span> <span class="o">=</span> <span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inputs&quot;</span><span class="p">,</span> <span class="p">[])</span>
        <span class="n">attr</span> <span class="o">=</span> <span class="n">reg_info</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;attr&quot;</span><span class="p">,</span> <span class="p">[])</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">tensor_inputs</span><span class="p">,</span> <span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">)):</span>
            <span class="n">tensor_inputs</span> <span class="o">=</span> <span class="p">[</span><span class="n">tensor_inputs</span><span class="p">]</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">attr</span><span class="p">,</span> <span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">)):</span>
            <span class="n">attr</span> <span class="o">=</span> <span class="p">[</span><span class="n">attr</span><span class="p">]</span>
        <span class="c1"># input_names include tensor input names and attr input names</span>
        <span class="n">input_names</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="c1"># attr_names only includes attr input names</span>
        <span class="n">attr_names</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">tensor_inputs</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">item</span><span class="p">,</span> <span class="nb">dict</span><span class="p">)</span> <span class="ow">and</span> <span class="n">item</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;name&quot;</span><span class="p">)</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">input_names</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">item</span><span class="p">[</span><span class="s2">&quot;name&quot;</span><span class="p">])</span>
        <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">attr</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">item</span><span class="p">,</span> <span class="nb">dict</span><span class="p">)</span> <span class="ow">and</span> <span class="n">item</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;name&quot;</span><span class="p">)</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">input_names</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">item</span><span class="p">[</span><span class="s2">&quot;name&quot;</span><span class="p">])</span>
                <span class="n">attr_names</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">item</span><span class="p">[</span><span class="s2">&quot;name&quot;</span><span class="p">])</span>
        <span class="n">cur_attr</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;input_names&quot;</span><span class="p">:</span> <span class="n">input_names</span><span class="p">,</span> <span class="s2">&quot;attr_names&quot;</span><span class="p">:</span> <span class="n">attr_names</span><span class="p">}</span>
        <span class="c1"># If func does not have attr, save current attr.</span>
        <span class="c1"># Else, check if current attr is same as previous saved one.</span>
        <span class="n">prev_input_names</span> <span class="o">=</span> <span class="n">input_names</span>
        <span class="n">prev_attr_names</span> <span class="o">=</span> <span class="n">attr_names</span>
        <span class="k">if</span> <span class="n">callable</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">):</span>
            <span class="n">func_attr</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="s2">&quot;func_attr&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">func_attr</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
                <span class="nb">setattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="s2">&quot;func_attr&quot;</span><span class="p">,</span> <span class="n">cur_attr</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">prev_input_names</span> <span class="o">=</span> <span class="n">func_attr</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;input_names&quot;</span><span class="p">)</span>
                <span class="n">prev_attr_names</span> <span class="o">=</span> <span class="n">func_attr</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;attr_names&quot;</span><span class="p">)</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="n">func_attr</span> <span class="o">=</span> <span class="n">Custom</span><span class="o">.</span><span class="n">attr_dict</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">)</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">func_attr</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
                <span class="n">Custom</span><span class="o">.</span><span class="n">attr_dict</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">]</span> <span class="o">=</span> <span class="n">cur_attr</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">prev_input_names</span> <span class="o">=</span> <span class="n">func_attr</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;input_names&quot;</span><span class="p">)</span>
                <span class="n">prev_attr_names</span> <span class="o">=</span> <span class="n">func_attr</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;attr_names&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">prev_input_names</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;func </span><span class="si">{}</span><span class="s2">: previous saved input_names should be a list, but got </span><span class="si">{}</span><span class="s2">&quot;</span>
                            <span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="nb">type</span><span class="p">(</span><span class="n">prev_input_names</span><span class="p">)))</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">input_names</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">prev_input_names</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;func </span><span class="si">{}</span><span class="s2">: input_names&#39;s length </span><span class="si">{}</span><span class="s2"> is different from previous saved one: </span><span class="si">{}</span><span class="s2">&quot;</span>
                             <span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">input_names</span><span class="p">),</span> <span class="nb">len</span><span class="p">(</span><span class="n">prev_input_names</span><span class="p">)))</span>
        <span class="k">if</span> <span class="n">attr_names</span> <span class="o">!=</span> <span class="n">prev_attr_names</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;func </span><span class="si">{}</span><span class="s2">: attr_names </span><span class="si">{}</span><span class="s2"> is different from previous saved one: </span><span class="si">{}</span><span class="s2">&quot;</span>
                             <span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="n">attr_names</span><span class="p">,</span> <span class="n">prev_attr_names</span><span class="p">))</span>

    <span class="k">def</span> <span class="nf">_update_attr</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Add input_names, attr_names, primitive_target to primitive&#39;s attr.&quot;&quot;&quot;</span>
        <span class="c1"># add input_names, attr_names</span>
        <span class="n">func_attr</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="k">if</span> <span class="n">callable</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">):</span>
            <span class="n">func_attr</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="s2">&quot;func_attr&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="n">func_attr</span> <span class="o">=</span> <span class="n">Custom</span><span class="o">.</span><span class="n">attr_dict</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">func_attr</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
            <span class="n">input_names</span> <span class="o">=</span> <span class="n">func_attr</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;input_names&quot;</span><span class="p">)</span>
            <span class="n">attr_names</span> <span class="o">=</span> <span class="n">func_attr</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;attr_names&quot;</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">input_names</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;input_names&quot;</span><span class="p">,</span> <span class="n">input_names</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">attr_names</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;attr_names&quot;</span><span class="p">,</span> <span class="n">attr_names</span><span class="p">)</span>
        <span class="c1"># add primitive_target</span>
        <span class="n">registered_targets</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_registered_targets</span><span class="p">()</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span> <span class="o">==</span> <span class="s2">&quot;pyfunc&quot;</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;primitive_target&quot;</span><span class="p">,</span> <span class="s2">&quot;CPU&quot;</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">registered_targets</span> <span class="ow">and</span> <span class="n">registered_targets</span> <span class="o">!=</span> <span class="p">[</span><span class="s2">&quot;CPU&quot;</span><span class="p">]:</span>
                <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;CustomPyfunc only supports CPU platform, but gets registered target as </span><span class="si">{}</span><span class="s2">.&quot;</span>
                               <span class="s2">&quot;We will run CustomPyfunc on CPU&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">registered_targets</span><span class="p">))</span>
        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">func_type</span> <span class="o">==</span> <span class="s2">&quot;aot&quot;</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">registered_targets</span><span class="p">)</span> <span class="o">!=</span> <span class="mi">1</span><span class="p">:</span>
                <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Target of CustomAOT will be set according to context.&quot;</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">registered_targets</span> <span class="o">==</span> <span class="p">[</span><span class="s2">&quot;GPU&quot;</span><span class="p">]:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;primitive_target&quot;</span><span class="p">,</span> <span class="s2">&quot;GPU&quot;</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">registered_targets</span> <span class="o">==</span> <span class="p">[</span><span class="s2">&quot;CPU&quot;</span><span class="p">]:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;primitive_target&quot;</span><span class="p">,</span> <span class="s2">&quot;CPU&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">callable</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">)</span> <span class="ow">and</span> <span class="n">callable</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">out_shape</span><span class="p">):</span>
            <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">out_shape</span><span class="p">,</span> <span class="s2">&quot;type&quot;</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">out_shape</span><span class="p">,</span> <span class="s2">&quot;type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;autodiff&quot;</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;autodiff&quot;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">add_prim_attr</span><span class="p">(</span><span class="s2">&quot;autodiff&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_hybrid_autodiff</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;generate backward op for a custom hybrid op&quot;&quot;&quot;</span>
        <span class="n">inputs_num</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">inspect</span><span class="o">.</span><span class="n">signature</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">)</span><span class="o">.</span><span class="n">parameters</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">inputs_num</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;Function with no input has no backward op.&quot;</span><span class="p">)</span>
        <span class="k">elif</span> <span class="n">inputs_num</span> <span class="o">&gt;</span> <span class="mi">10</span><span class="p">:</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;Currently autodiff for function with more than 10 inputs is not supported.&quot;</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">grad_func</span> <span class="o">=</span> <span class="n">autodiff_bprop</span><span class="p">(</span><span class="n">inputs_num</span><span class="p">)</span>

            <span class="k">def</span> <span class="nf">infer_func</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">):</span>
                <span class="k">return</span> <span class="n">args</span><span class="p">[:</span><span class="n">inputs_num</span><span class="p">]</span>
            <span class="nb">setattr</span><span class="p">(</span><span class="n">infer_func</span><span class="p">,</span> <span class="s2">&quot;type&quot;</span><span class="p">,</span> <span class="s2">&quot;autodiff&quot;</span><span class="p">)</span>
            <span class="n">op</span> <span class="o">=</span> <span class="n">Custom</span><span class="p">(</span><span class="n">func</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">,</span> <span class="n">out_shape</span><span class="o">=</span><span class="n">infer_func</span><span class="p">,</span> <span class="n">out_dtype</span><span class="o">=</span><span class="n">infer_func</span><span class="p">,</span>
                        <span class="n">func_type</span><span class="o">=</span><span class="s2">&quot;akg&quot;</span><span class="p">,</span> <span class="n">bprop</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">bprop</span> <span class="o">=</span> <span class="n">grad_func</span><span class="p">(</span><span class="n">op</span><span class="p">)</span></div>
</pre></div>

           </div>
           
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>
        &#169; Copyright 2021, MindSpore.

    </p>
  </div>
    
    
    
    Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    
    provided by <a href="https://readthedocs.org">Read the Docs</a>. 

</footer>
        </div>
      </div>

    </section>

  </div>
  

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
   

</body>
</html>